"""Pellet model"""

import numpy as N
import bvp
import phyconst


class Pellet():
    """Pellet model"""
    def __init__(self, ratelaw, r, d, diffArray, porosity, tortuosity):
        """ratelaw: Rate law object
        r: pellet Radius in cm
        d: pellet Density in g/ml
        diffArray: component effective Diffusivity array in cm**2/s
        porosity: pellet inner Porosity
        tortuosity: pellet inner Tortuosity      
        """
        self.ratelaw = ratelaw
        self.r = r
        self.d = d
        self.diffArray = diffArray
        self.porosity = porosity
        self.tortuosity = tortuosity
        
    def calc(self, hasDiff, p, propArray, t, shellNum, carbonNum):
        """Calculate the reaction rate, selectivity, effective factor
        and the reaction heat(delta H)
        hasDiff: Has diffusion?
        p: pellet surface Pressure in Pa
        propArray: component pressure Proportion array [pCO, pH2, pCH4, pH2O, [pCO2], [pC2+]]
        t: pellet surface Temperature in K
        shellNum: Shell number in solving diffusion
        """
        self.t = t
        self.componentNum = len(propArray)
        
        if propArray[0] == 0.0 or propArray[1] == 0.0:
            raArray = N.zeros((self.componentNum))
            effCO = 0.0
            selCH4 = 0.0
            sels = N.array([selCH4])
            if carbonNum > 1:
                selCO2 = 0.0
                sels = N.r_[selCH4, selCO2, N.zeros(2 * (carbonNum - 1))]
            deltaH = 0.0
            profileArray = N.zeros((self.componentNum, shellNum))
            return raArray, effCO, sels, deltaH, profileArray
        
        pArray = p * propArray
        raArrayS = self.ratelaw.calc_rate(pArray, t)
        if not hasDiff:
            raArray = raArrayS
            effCO = 1.0
            selCH4 = raArray[2] / -raArray[0]
            sels = N.array([selCH4])
            if carbonNum > 1:
                selCO2 = raArray[4] / -raArray[0]
                mols = N.arange(2.0, carbonNum + 1)
                mols = N.r_[mols, mols]    # paraffin and olefin
                c2plusArray = mols * raArray[5:] / -raArray[0]
                sels = N.r_[selCH4, selCO2, c2plusArray]
            deltaH = self.ratelaw.calc_heat(raArray, sels)
            profileArray = N.zeros((self.componentNum, shellNum))
            return raArray, effCO, sels, deltaH, profileArray
        
        self.cArrayS = pArray / phyconst.Rgas / t * 1e-6   # components' concentration on surface 
        boundaryPoints = N.r_[N.zeros(self.componentNum), N.ones(self.componentNum) * self.r]
        tol = N.ones(2 * self.componentNum) * 1e-5   # match z vector [u1, du1, u2, du2...]
        tol[1::2] = 0.0
        degrees = N.ones(self.componentNum) * 2   # degrees[i] is the ith equation's degree
        
        solution = bvp.colnew.solve(boundaryPoints, degrees, self.fsub, self.gsub,
                                    tolerances=tol, 
                                    vectorized=True, maximum_mesh_size=300,
                                    initial_guess = self.guess)
        
        x = N.linspace(0.0, self.r, shellNum)
        s = solution(x)
        shellThickness = self.r / shellNum
        raProfile = N.array([self.ratelaw.calc_rate(cArray * phyconst.Rgas * self.t  * 1e6, t) for cArray in s[:, ::2]])
        raProfile = raProfile[1::, :]  # Remove r = 0.0 value  
        shellAreaArray = 4.0 * N.pi * x[1::]**2 * shellThickness
        raProfile = shellAreaArray * raProfile.T
        raArray = N.array([N.sum(raItemProfile) for raItemProfile in raProfile]) / (4.0 / 3.0 * N.pi * self.r**3)
        effCO = raArray[0] / raArrayS[0]
        if effCO > 1.0:
            effCO = 1.0

        selCH4 = raArray[2] / -raArray[0]
        sels = N.array([selCH4])
        if carbonNum > 1:
            selCO2 = raArray[4] / -raArray[0]
            mols = N.arange(2.0, carbonNum + 1)
            mols = N.r_[mols, mols]    # paraffin and olefin
            c2plusArray = mols * raArray[5:] / -raArray[0]
            sels = N.r_[selCH4, selCO2, c2plusArray]
        deltaH = self.ratelaw.calc_heat(raArray, sels)
        profileArray = s[:, ::2].T
        return raArray, effCO, sels, deltaH, profileArray
      
        
    def fsub(self, x, z):
        """The equations: C" = - 2 / r * C' - density / diff * ra"""
        du = z[1::2]   #u1, du1, u2, du2 = z  it's neat to name the variables
        u = z[::2]
        pArray = u * phyconst.Rgas * self.t  * 1e6  # (componentNum, meshNum)
        raArray = N.zeros(pArray.shape)
        for mesh in range(pArray.shape[1]):
            raArray[:, mesh] = self.ratelaw.calc_rate(pArray[:, mesh], self.t)
        zeroIndies = N.where(x == 0.0)[0]  # 'where' returns a tuple
        if not zeroIndies:
            return -2.0 / x * du - (self.d / self.diffArray * raArray.T).T
        else:
            nonzeroIndies = x.nonzero()
            ddy = N.zeros(u.shape)
            ddy[:, nonzeroIndies] = -2.0 / x[nonzeroIndies] * du[:, nonzeroIndies] - \
                (self.d / self.diffArray * raArray[:, nonzeroIndies].T).T
            ddy[:, zeroIndies] = -(self.d / self.diffArray * raArray[:, nonzeroIndies].T).T / 3.0
            return ddy
          
    
    def gsub(self, z):
        """The boundary conditions"""
        bPoints = z.shape[1]
        du = z[1::2]
        u = z[::2]
        du_ = du[range(0, bPoints/2), range(0, bPoints/2)]  # [du[0, 0], du[1, 1], du[2, 2],...]
        u_ = u[range(0, bPoints/2), range(bPoints/2, bPoints)]  # [u[0, 4], u[1, 5], u[2, 6], ...]
        return N.r_[du_, u_ - self.cArrayS]
    
    def guess(self, x):
        """Initial guess"""
        n = self.componentNum * 2
        z = N.zeros((n,) + N.asarray(x).shape)
        dm = N.zeros((n / 2,) + N.asarray(x).shape)
        for comp in range(0, self.componentNum):
            z[comp * 2] = self.cArrayS[comp]
        return z, dm

    
            
            

